import os

import torch
import torch.nn as nn
from torch import optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

from vit import VisionTransformer


def main():
    # 数据加载
    transform = transforms.Compose([
        transforms.Resize((224, 224)),  # 或考虑其他方法处理小图像
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

    train_set = datasets.CIFAR100(root="../Resnet/data", train=True, download=True, transform=transform)
    train_loader = DataLoader(train_set, batch_size=64, shuffle=True)

    test_set = datasets.CIFAR100(root="../Resnet/data", train=False, download=True, transform=transform)
    test_loader = DataLoader(test_set, batch_size=64, shuffle=False)

    # 初始化模型
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = VisionTransformer(img_size=224, patch_size=16, in_channels=3, embed_size=768, num_patches=196, heads=4,
                              forward_expansion=4, layers=12, dropout=0.1, num_classes=100).to(device)

    # 训练配置
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

    # 训练模型并每3轮保存一次模型
    def train_model(num_epochs):
        for epoch in range(num_epochs):
            model.train()
            print('Epoch {}/{}'.format(epoch, num_epochs), end=' ')
            total_loss = 0
            for images, labels in train_loader:
                images, labels = images.to(device), labels.to(device)

                optimizer.zero_grad()
                outputs = model(images)
                loss = criterion(outputs, labels)
                loss.backward()
                optimizer.step()
                total_loss += loss.item()

            print(f"Epoch {epoch + 1}, Loss: {total_loss / len(train_loader)}")

            # 每3轮保存一次模型
            if (epoch + 1) % 3 == 0:
                save_path = f"./model_checkpoints/vision_transformer_cifar100_epoch_{epoch + 1}.pt"
                os.makedirs(os.path.dirname(save_path), exist_ok=True)
                torch.save(model.state_dict(), save_path)
                print(f"Model saved after epoch {epoch + 1}")

    train_model(60)

    # 测试模型
    def test_model():
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for images, labels in test_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
        print(f'Accuracy of the model on the {total} test images: {100 * correct / total}%')

    test_model()


if __name__ == "__main__":
    main()
    # Accuracy of the model on the 10000 test images: 44.01%
